{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# for real data\n",
    "# create dataset of number sequences\n",
    "# let's assume that we have a vocabulary size of 1000 words\n",
    "# let's assume that 0 is the EOS token, and 1 is the SOS token, and 2 is PAD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "Epoch  1\n",
      "[100/980] Loss: 0.3453\n",
      "[200/980] Loss: 0.3516\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-037865ba84cc>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    121\u001b[0m                 tmp_out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n\u001b[1;32m    122\u001b[0m                                 \u001b[0mencoded_idx\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minput_out\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprev_state\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 123\u001b[0;31m                                 weighted=w, order=j)\n\u001b[0m\u001b[1;32m    124\u001b[0m                 \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtmp_out\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/home/mjc/anaconda3/envs/pytorch/lib/python3.5/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    204\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    205\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 206\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    207\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    208\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/home/mjc/projects/copynet/models/copynet_debug.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_idx, encoded, encoded_idx, prev_state, weighted, order)\u001b[0m\n\u001b[1;32m    124\u001b[0m                 \u001b[0midx_from_input\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m                 \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mj\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mencoded_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 126\u001b[0;31m                         \u001b[0midx_from_input\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0minput_idx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mj\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    127\u001b[0m                 \u001b[0midx_from_input\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx_from_input\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    128\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mI\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx_from_input\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/home/mjc/projects/copynet/models/copynet_debug.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    124\u001b[0m                 \u001b[0midx_from_input\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m                 \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mj\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mencoded_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 126\u001b[0;31m                         \u001b[0midx_from_input\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0minput_idx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mj\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    127\u001b[0m                 \u001b[0midx_from_input\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx_from_input\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    128\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mI\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx_from_input\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "from models.copynet_debug import CopyEncoder, CopyDecoder\n",
    "from models.functions import numpy_to_var, toData, to_np, to_var, visualize, decoder_initial, update_logger\n",
    "import time\n",
    "import sys\n",
    "import math\n",
    "torch.manual_seed(1000)\n",
    "\n",
    "# Hyperparameters\n",
    "embed_size = 150\n",
    "hidden_size = 300\n",
    "num_layers = 1\n",
    "bin_size = 10\n",
    "num_epochs = 40\n",
    "prev_end=0\n",
    "batch_size = 100\n",
    "lr = 0.001\n",
    "vocab_size = 100\n",
    "weight_decay = 0.99\n",
    "use_saved = True # whether to train from a previous model\n",
    "continue_from = 70\n",
    "step = 0 # number of steps taken\n",
    "\n",
    "# input and output directories\n",
    "version = 'v6'\n",
    "vocab_size = 1000\n",
    "data_dir = 'data/copynet_data_%s.txt' %(version)\n",
    "save_dir = 'results/copynet_data_%s.csv' %(version)\n",
    "\n",
    "\n",
    "# get training and test data\n",
    "with open(data_dir) as f:\n",
    "    lines = f.readlines()\n",
    "import random\n",
    "half = int(len(lines)/2)\n",
    "train = lines\n",
    "random.shuffle(train)\n",
    "test = train[:20]\n",
    "train = train[20:]\n",
    "\n",
    "# get number of batches\n",
    "num_samples = len(train)\n",
    "num_batches = int(num_samples/batch_size)\n",
    "\n",
    "\n",
    "################ load copynet model #####################\n",
    "if use_saved:\n",
    "    # if using from previous data\n",
    "    encoder_dir = 'models/encoder_%s_%s.pckl' % (version,continue_from)\n",
    "    decoder_dir = 'models/decoder_%s_%s.pckl' % (version,continue_from)\n",
    "    encoder = torch.load(f=encoder_dir)\n",
    "    decoder = torch.load(f=decoder_dir)\n",
    "else:\n",
    "    encoder = CopyEncoder(vocab_size, embed_size, hidden_size)\n",
    "    decoder = CopyDecoder(vocab_size, embed_size, hidden_size)\n",
    "    continue_from = 0\n",
    "if torch.cuda.is_available():\n",
    "    encoder.cuda()\n",
    "    decoder.cuda()\n",
    "\n",
    "################################# training ##################################\n",
    "\n",
    "# set loss\n",
    "criterion = nn.NLLLoss()\n",
    "\n",
    "start = time.time()\n",
    "for epoch in range(num_epochs):\n",
    "    print(\"==================================================\")\n",
    "    print(\"Epoch \",epoch+1)\n",
    "    opt_e = optim.Adam(params=encoder.parameters(), lr=lr)\n",
    "    opt_d = optim.Adam(params=decoder.parameters(), lr=lr)\n",
    "    lr= lr * weight_decay # weight decay\n",
    "    # shuffle data\n",
    "    random.shuffle(train)\n",
    "    samples_read = 0\n",
    "    while(samples_read<len(train)):\n",
    "        # initialize gradient buffers\n",
    "        opt_e.zero_grad()\n",
    "        opt_d.zero_grad()\n",
    "\n",
    "        # obtain batch outputs\n",
    "        batch = train[samples_read:min(samples_read+batch_size,len(train))]\n",
    "        input_out, output_out, in_len, out_len = toData(batch)\n",
    "        samples_read+=len(batch)\n",
    "\n",
    "        # mask input to remove padding\n",
    "        input_mask = np.array(input_out>0, dtype=int)\n",
    "\n",
    "        # input and output in Variable form\n",
    "        x = numpy_to_var(input_out)\n",
    "        y = numpy_to_var(output_out)\n",
    "\n",
    "        # apply to encoder\n",
    "        encoded, _ = encoder(x)\n",
    "\n",
    "        # get initial input of decoder\n",
    "        decoder_in, s, w = decoder_initial(x.size(0))\n",
    "\n",
    "        # out_list to store outputs\n",
    "        out_list=[]\n",
    "        for j in range(y.size(1)): # for all sequences\n",
    "            \"\"\"\n",
    "            decoder_in (Variable): [b]\n",
    "            encoded (Variable): [b x seq x hid]\n",
    "            input_out (np.array): [b x seq]\n",
    "            s (Variable): [b x hid]\n",
    "            \"\"\"\n",
    "            # 1st state\n",
    "            if j==0:\n",
    "                out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n",
    "                                encoded_idx=input_out, prev_state=s,\n",
    "                                weighted=w, order=j)\n",
    "            # remaining states\n",
    "            else:\n",
    "                tmp_out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n",
    "                                encoded_idx=input_out, prev_state=s,\n",
    "                                weighted=w, order=j)\n",
    "                out = torch.cat([out,tmp_out],dim=1)\n",
    "\n",
    "            # for debugging: stop if nan\n",
    "            if math.isnan(w[-1][0][0].data[0]):\n",
    "                sys.exit()\n",
    "            # select next input\n",
    "\n",
    "            if epoch % 2 ==13:\n",
    "                decoder_in = out[:,-1].max(1)[1].squeeze() # train with sequence outputs\n",
    "            else:\n",
    "                decoder_in = y[:,j] # train with ground truth\n",
    "            out_list.append(out[:,-1].max(1)[1].squeeze().cpu().data.numpy())\n",
    "\n",
    "        # print(torch.stack(decoder.prob_c_to_g,1))\n",
    "        target = pack_padded_sequence(y,out_len.tolist(), batch_first=True)[0]\n",
    "        pad_out = pack_padded_sequence(out,out_len.tolist(), batch_first=True)[0]\n",
    "        # include log computation as we are using log-softmax and NLL\n",
    "        pad_out = torch.log(pad_out)\n",
    "        loss = criterion(pad_out, target)\n",
    "        loss.backward()\n",
    "        # if samples_read%500==0:\n",
    "        print(\"[%d/%d] Loss: %1.4f\"%(samples_read,len(train),loss.data[0]))\n",
    "        opt_e.step()\n",
    "        opt_d.step()\n",
    "        step += 1\n",
    "        info = {\n",
    "            'loss': loss.data[0]\n",
    "        }\n",
    "    # print(\"Loss: \",loss.data[0])\n",
    "    elapsed = time.time()\n",
    "    print(\"Elapsed time for epoch: \",elapsed-start)\n",
    "    start = time.time()\n",
    "\n",
    "    ################################# validation ##################################\n",
    "    if epoch % 10==0:\n",
    "        print(\"Printing results\")\n",
    "        input_out, output_out, in_len, out_len = toData(test)\n",
    "        input_mask = np.array(input_out>0, dtype=int)\n",
    "        x = numpy_to_var(input_out)\n",
    "        y = numpy_to_var(output_out)\n",
    "        if epoch==0:\n",
    "            with open(save_dir,'a') as f:\n",
    "                x_input = x.cpu().data.numpy().transpose()\n",
    "                y_input = y.cpu().data.numpy().transpose()\n",
    "                m = max(x_input.shape[0],y_input.shape[0])\n",
    "                out_ = np.zeros([m,x_input.shape[1]*2])\n",
    "                out_[:x_input.shape[0],x_input.shape[1]:]=x_input\n",
    "                out_[:y_input.shape[0],:y_input.shape[1]]=y_input\n",
    "                out = []\n",
    "                for line in out_:\n",
    "                    tmp = ','.join([str(x_) for x_ in line])\n",
    "                    out.append(tmp)\n",
    "                f.write('\\n'.join(out)+'\\n')\n",
    "        encoded, _ = encoder(x)\n",
    "        decoder_in, s, w = decoder_initial(x.size(0))\n",
    "        out_list=[]\n",
    "        for j in range(y.size(1)): # for all sequences\n",
    "            if j==0:\n",
    "                out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n",
    "                                encoded_idx=input_out, prev_state=s,\n",
    "                                weighted=w, order=j)\n",
    "            else:\n",
    "                tmp_out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n",
    "                                encoded_idx=input_out, prev_state=s,\n",
    "                                weighted=w, order=j)\n",
    "                out = torch.cat([out,tmp_out],dim=1)\n",
    "            decoder_in = out[:,-1].max(1)[1].squeeze() # train with sequence outputs\n",
    "            out_list.append(out[:,-1].max(1)[1].squeeze().cpu().data.numpy())\n",
    "        out = np.array(out_list)\n",
    "        iden = np.array(out==y_input,dtype=int)\n",
    "        with open(save_dir,'a') as f:\n",
    "            out = np.hstack(tup=(out,iden))\n",
    "            f.write('\\n')\n",
    "            for line in out:\n",
    "                f.write(','.join([str(y_) for y_ in line])+'\\n')\n",
    "        torch.save(f='models/encoder_%s_%s.pckl' % (version,str(epoch+continue_from)),obj=encoder)\n",
    "        torch.save(f='models/decoder_%s_%s.pckl' % (version,str(epoch+continue_from)),obj=decoder)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Printing results\n"
     ]
    }
   ],
   "source": [
    "################################# validation ##################################\n",
    "print(\"Printing results\")\n",
    "input_out, output_out, in_len, out_len = toData(test[:10])\n",
    "input_mask = np.array(input_out>0, dtype=int)\n",
    "x = numpy_to_var(input_out)\n",
    "y = numpy_to_var(output_out)\n",
    "encoded, _ = encoder(x)\n",
    "decoder_in, s, w = decoder_initial(x.size(0))\n",
    "out_list=[]\n",
    "for j in range(y.size(1)): # for all sequences\n",
    "    if j==0:\n",
    "        out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n",
    "                        encoded_idx=input_out, prev_state=s, \n",
    "                        weighted=w, order=j)\n",
    "    else:\n",
    "        tmp_out, s, w = decoder(input_idx=decoder_in, encoded=encoded,\n",
    "                        encoded_idx=input_out, prev_state=s, \n",
    "                        weighted=w, order=j)\n",
    "        out = torch.cat([out,tmp_out],dim=1)\n",
    "    decoder_in = y[:,j] # train with ground truth\n",
    "#     decoder_in = out[:,-1].max(1)[1].squeeze() # train with sequence outputs\n",
    "    out_list.append(out[:,-1].max(1)[1].squeeze().cpu().data.numpy())\n",
    "# with open(save_dir,'a') as f:\n",
    "#     out = np.hstack(tup=(out,iden))\n",
    "#     f.write('\\n')\n",
    "#     for line in out:\n",
    "#         f.write(','.join([str(y_) for y_ in line])+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input:  [502 840 679 776 209 707 471 490 267   1]\n",
      "truth:  [467 864 939 103 840 679 776 209 863 938   1]\n",
      "output:  [467 864 939 103 840 840 776 776 863 938   1]\n"
     ]
    }
   ],
   "source": [
    "# debug\n",
    "# get a sample input, ground truth, output\n",
    "idx = 0\n",
    "print(\"input: \",x[idx].cpu().data.numpy())\n",
    "print(\"truth: \",y[idx].cpu().data.numpy())\n",
    "O = torch.cat(decoder.O,1)\n",
    "print(\"output: \",O[idx].max(1)[1].cpu().numpy().squeeze())\n",
    "A = torch.stack(decoder.A,1)\n",
    "A2 = torch.stack(decoder.A2,1)\n",
    "P = torch.stack(decoder.P,1)\n",
    "I = torch.stack(decoder.I,1)\n",
    "E = torch.stack(decoder.E,1)\n",
    "S = torch.stack(decoder.S,1)\n",
    "W = torch.cat(decoder.W,1)\n",
    "Y = torch.stack(decoder.Y,1)\n",
    "scores = torch.stack(decoder.scores,1)\n",
    "sc = torch.stack(decoder.sc,1)\n",
    "pcg = torch.stack(decoder.prob_c_to_g,1)\n",
    "pc = P[:,:,1000:]\n",
    "pg = P[:,:,:1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYcAAAD8CAYAAACcjGjIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAEktJREFUeJzt3X+s3fVdx/Hny9uWDYQhwyG2xVXTmHUqjDUFN+LAuVlw\ns5rsjzb7ocvIzcxI5o9oOk22qH/pzDQLuOZmqzgdELPB1ixlDHSROcLWMhFaoNtdwdErWjcm7FcK\n1bd/nG/leD/39pzbntt76H0+kpP7/X5+fM/n+6Hti8/3+z3npqqQJKnfDyz1ACRJ48dwkCQ1DAdJ\nUsNwkCQ1DAdJUsNwkCQ1BoZDkrVJPpfkoST7k7x7jjZJ8sEk00keSHJpX93mJAe6uu2jPgFJ0ugN\ns3I4CvxOVW0ALgfelWTDrDZXA+u71yTwIYAkE8ANXf0GYNscfSVJY2ZgOFTVE1X15W7728DDwOpZ\nzbYAH62ee4Fzk1wIbAKmq+pgVT0D3NK1lSSNsRULaZzkpcArgC/OqloNPN63f6grm6v8snmOPUlv\n1cEEE688k3MWMjQtQM44gzpyZKmHoTH3wpfB9x9e6lGM3ul6Xt/mW9+oqh8e1fGGDockPwh8AvjN\nqnp6VAM4pqqmgCmAc3JeXZbXjvot1FnxYz/B0a9+bamHoTG34WMreOiVR5d6GCN3up7XXfXxfx3l\n8YYKhyQr6QXDx6rq1jmazABr+/bXdGUr5ymXJI2xYZ5WCvAR4OGq+sA8zXYBb+ueWroceKqqngD2\nAOuTrEuyCtjatZUkjbFhVg6vBt4KPJjk/q7s94GLAKpqB7AbuAaYBr4HvL2rO5rkOuAOYALYWVX7\nR3oGkqSRGxgOVfVPQAa0KeBd89TtphcekqTnCT8hLUlqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbh\nIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpMbAXxOa\nZCfwBuBwVf3UHPW/C7y573gvA364qp5M8hjwbeC/gaNVtXFUA5ckLZ5hVg43Apvnq6yq91fVJVV1\nCfAe4B+r6sm+Jld19QaDJD1PDAyHqrobeHJQu8424OaTGpEkacmN7J5DkjPprTA+0VdcwF1J7ksy\nOar3kiQtroH3HBbgjcAXZl1SuqKqZpK8BLgzySPdSqTRhcckwAs4c4TDkiQt1CifVtrKrEtKVTXT\n/TwM3AZsmq9zVU1V1caq2riSM0Y4LEnSQo0kHJK8CHgN8Km+srOSnH1sG3g9sG8U7ydJWlzDPMp6\nM3AlcH6SQ8D7gJUAVbWja/arwGer6rt9XS8Abkty7H1uqqrPjG7okqTFMjAcqmrbEG1upPfIa3/Z\nQeDiEx2YJGnp+AlpSVLDcJAkNQwHSVLDcJAkNQwHSVLDcJAkNQwHSVLDcJAkNQwHSVLDcJAkNQwH\nSVLDcJAkNQwHSVLDcJAkNQwHSVLDcJAkNQwHSVLDcJAkNQwHSVJjYDgk2ZnkcJJ989RfmeSpJPd3\nr/f21W1OciDJdJLtoxy4JGnxDLNyuBHYPKDN56vqku71RwBJJoAbgKuBDcC2JBtOZrCSpFNjYDhU\n1d3Akydw7E3AdFUdrKpngFuALSdwHEnSKTaqew6vSvJAktuTvLwrWw083tfmUFc2pySTSfYm2fss\nR0Y0LEnSiVgxgmN8Gbioqr6T5Brgk8D6hR6kqqaAKYBzcl6NYFySpBN00iuHqnq6qr7Tbe8GViY5\nH5gB1vY1XdOVSZLG3EmHQ5IfSZJue1N3zG8Ce4D1SdYlWQVsBXad7PtJkhbfwMtKSW4GrgTOT3II\neB+wEqCqdgBvAn4jyVHg+8DWqirgaJLrgDuACWBnVe1flLOQJI3UwHCoqm0D6q8Hrp+nbjew+8SG\nJklaKn5CWpLUMBwkSQ3DQZLUMBwkSQ3DQZLUMBwkSQ3DQZLUMBwkSQ3DQZLUMBwkSQ3DQZLUMBwk\nSQ3DQZLUMBwkSQ3DQZLUMBwkSQ3DQZLUMBwkSY2B4ZBkZ5LDSfbNU//mJA8keTDJPUku7qt7rCu/\nP8neUQ5ckrR4hlk53AhsPk79o8BrquqngT8GpmbVX1VVl1TVxhMboiTpVFsxqEFV3Z3kpcepv6dv\n915gzckPS5K0lEZ9z+EdwO19+wXcleS+JJPH65hkMsneJHuf5ciIhyVJWoiBK4dhJbmKXjhc0Vd8\nRVXNJHkJcGeSR6rq7rn6V9UU3SWpc3JejWpckqSFG8nKIcnPAB8GtlTVN4+VV9VM9/MwcBuwaRTv\nJ0laXCcdDkkuAm4F3lpVX+krPyvJ2ce2gdcDcz7xJEkaLwMvKyW5GbgSOD/JIeB9wEqAqtoBvBd4\nMfCXSQCOdk8mXQDc1pWtAG6qqs8swjlIkkZsmKeVtg2ovxa4do7yg8DFbQ9J0rjzE9KSpIbhIElq\nGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6SpIbhIElqGA6S\npIbhIElqGA6SpIbhIElqDAyHJDuTHE6yb576JPlgkukkDyS5tK9uc5IDXd32UQ5ckrR4hlk53Ahs\nPk791cD67jUJfAggyQRwQ1e/AdiWZMPJDFaSdGoMDIequht48jhNtgAfrZ57gXOTXAhsAqar6mBV\nPQPc0rWVJI25UdxzWA083rd/qCubr3xOSSaT7E2y91mOjGBYkqQTNTY3pKtqqqo2VtXGlZyx1MM5\nrdULVy31EPQ8cODt65d6CIvizy+8b6mH8LywYgTHmAHW9u2v6cpWzlMuSRpzo1g57ALe1j21dDnw\nVFU9AewB1idZl2QVsLVrK0kacwNXDkluBq4Ezk9yCHgfvVUBVbUD2A1cA0wD3wPe3tUdTXIdcAcw\nAeysqv2LcA6SpBEbGA5VtW1AfQHvmqduN73wkCQ9j4zNDWlJ0vgwHCRJDcNBktQwHCRJDcNBktQw\nHCRJDcNBktQwHCRJDcNBktQwHCRJDcNBktQwHCRJDcNBktQwHCRJDcNBktQwHCRJDcNBktQwHCRJ\njaHCIcnmJAeSTCfZPkf97ya5v3vtS/LfSc7r6h5L8mBXt3fUJyBJGr2Bv0M6yQRwA/A64BCwJ8mu\nqnroWJuqej/w/q79G4Hfqqon+w5zVVV9Y6QjlyQtmmFWDpuA6ao6WFXPALcAW47Tfhtw8ygGJ0la\nGsOEw2rg8b79Q11ZI8mZwGbgE33FBdyV5L4kk/O9SZLJJHuT7H2WI0MMS5K0WAZeVlqgNwJfmHVJ\n6YqqmknyEuDOJI9U1d2zO1bVFDAFcE7OqxGPS5K0AMOsHGaAtX37a7qyuWxl1iWlqprpfh4GbqN3\nmUqSNMaGCYc9wPok65KsohcAu2Y3SvIi4DXAp/rKzkpy9rFt4PXAvlEMXJK0eAZeVqqqo0muA+4A\nJoCdVbU/yTu7+h1d018FPltV3+3rfgFwW5Jj73VTVX1mlCcgSRq9oe45VNVuYPessh2z9m8EbpxV\ndhC4+KRGKEk65fyEtCSpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySpYThIkhqG\ngySpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySpYThIkhpDhUOSzUkOJJlOsn2O+iuTPJXk/u71\n3mH7SpLGz8DfIZ1kArgBeB1wCNiTZFdVPTSr6eer6g0n2FeSNEaGWTlsAqar6mBVPQPcAmwZ8vgn\n01eStESGCYfVwON9+4e6stleleSBJLcnefkC+5JkMsneJHuf5cgQw5IkLZaBl5WG9GXgoqr6TpJr\ngE8C6xdygKqaAqYAzsl5NaJxSZJOwDArhxlgbd/+mq7s/1TV01X1nW57N7AyyfnD9JUkjZ9hwmEP\nsD7JuiSrgK3Arv4GSX4kSbrtTd1xvzlMX0nS+Bl4Wamqjia5DrgDmAB2VtX+JO/s6ncAbwJ+I8lR\n4PvA1qoqYM6+i3QukqQRGeqeQ3epaPessh1929cD1w/bV5I03vyEtCSpYThIkhqGgySpYThIkhqG\ngySpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySpYThIkhqGgySp\nYThIkhpDhUOSzUkOJJlOsn2O+jcneSDJg0nuSXJxX91jXfn9SfaOcvCSpMUx8HdIJ5kAbgBeBxwC\n9iTZVVUP9TV7FHhNVX0rydXAFHBZX/1VVfWNEY5bkrSIhlk5bAKmq+pgVT0D3AJs6W9QVfdU1be6\n3XuBNaMdpiTpVBomHFYDj/ftH+rK5vMO4Pa+/QLuSnJfksmFD1GSdKoNvKy0EEmuohcOV/QVX1FV\nM0leAtyZ5JGqunuOvpPAJMALOHOUw5IkLdAwK4cZYG3f/pqu7P9J8jPAh4EtVfXNY+VVNdP9PAzc\nRu8yVaOqpqpqY1VtXMkZw5+BJGnkhgmHPcD6JOuSrAK2Arv6GyS5CLgVeGtVfaWv/KwkZx/bBl4P\n7BvV4CVJi2PgZaWqOprkOuAOYALYWVX7k7yzq98BvBd4MfCXSQCOVtVG4ALgtq5sBXBTVX1mUc5E\nkjQyQ91zqKrdwO5ZZTv6tq8Frp2j30Hg4tnlkqTx5iekJUkNw0GS1DAcJEkNw0GS1DAcJEkNw0GS\n1DAcJEkNw0GS1DAcJEkNw0GS1DAcJEkNw0GS1DAcJEkNw0GS1DAcJEkNw0GS1DAcJEkNw0GS1DAc\nJEmNocIhyeYkB5JMJ9k+R32SfLCrfyDJpcP2lSSNn4HhkGQCuAG4GtgAbEuyYVazq4H13WsS+NAC\n+kqSxswwK4dNwHRVHayqZ4BbgC2z2mwBPlo99wLnJrlwyL6SpDGzYog2q4HH+/YPAZcN0Wb1kH0B\nSDJJb9UBcOSu+vi+Ica2HJwPfGOkR/yXkR7tVBn9PDx/nZq5eH78OVnwXExcCPDVRRnMEvvJUR5s\nmHA4JapqCpgCSLK3qjYu8ZDGgnPR4zw8x7l4jnPxnCR7R3m8YcJhBljbt7+mKxumzcoh+kqSxsww\n9xz2AOuTrEuyCtgK7JrVZhfwtu6ppcuBp6rqiSH7SpLGzMCVQ1UdTXIdcAcwAeysqv1J3tnV7wB2\nA9cA08D3gLcfr+8Q45o6kZM5TTkXPc7Dc5yL5zgXzxnpXKSqRnk8SdJpwE9IS5IahoMkqTFW4bDc\nvmojydokn0vyUJL9Sd7dlZ+X5M4kX+1+/lBfn/d083MgyS8u3ehHL8lEkn9O8uluf1nOA0CSc5N8\nPMkjSR5O8rPLcT6S/Fb3d2NfkpuTvGA5zUOSnUkOJ9nXV7bg80/yyiQPdnUfTJKBb15VY/Gid8P6\na8CPA6vofQRnw1KPa5HP+ULg0m77bOAr9L5m5E+B7V35duBPuu0N3bycAazr5mtiqc9jhPPx28BN\nwKe7/WU5D905/jVwbbe9Cjh3uc0HvQ/RPgq8sNv/O+DXl9M8AD8HXArs6ytb8PkDXwIuBwLcDlw9\n6L3HaeWw7L5qo6qeqKovd9vfBh6m9xdiC71/HOh+/kq3vQW4paqOVNWj9J4O23RqR704kqwBfgn4\ncF/xspsHgCQvovePwkcAquqZqvovlud8rABemGQFcCbwbyyjeaiqu4EnZxUv6Py7rzI6p6rurV5S\nfLSvz7zGKRzm+wqOZSHJS4FXAF8ELqje50QA/h24oNs+nefoL4DfA/6nr2w5zgP0/q/vP4G/6i6z\nfTjJWSyz+aiqGeDPgK8DT9D7/NRnWWbzMIeFnv/qbnt2+XGNUzgsW0l+EPgE8JtV9XR/XZf0p/Xz\nxkneAByuqvvma7Mc5qHPCnqXEj5UVa8Avkvv8sH/WQ7z0V1L30IvLH8UOCvJW/rbLId5OJ7FPP9x\nCodhvqbjtJNkJb1g+FhV3doV/0e3FKT7ebgrP13n6NXALyd5jN7lxJ9P8rcsv3k45hBwqKq+2O1/\nnF5YLLf5+AXg0ar6z6p6FrgVeBXLbx5mW+j5z3Tbs8uPa5zCYdl91Ub3xMBHgIer6gN9VbuAX+u2\nfw34VF/51iRnJFlH7/dnfOlUjXexVNV7qmpNVb2U3n/3f6iqt7DM5uGYqvp34PEkx75l87XAQyy/\n+fg6cHmSM7u/K6+ld19uuc3DbAs6/+4S1NNJLu/m8W19fea31HfjZ92Zv4beEztfA/5gqcdzCs73\nCnpLwgeA+7vXNcCLgb+n973CdwHn9fX5g25+DjDEEwfPtxdwJc89rbSc5+ESYG/3Z+OTwA8tx/kA\n/hB4BNgH/A29J3GWzTwAN9O73/IsvRXlO07k/IGN3Rx+Dbie7tsxjvfy6zMkSY1xuqwkSRoThoMk\nqWE4SJIahoMkqWE4SJIahoMkqWE4SJIa/wsakj4XYqzsVAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7f139869edd8>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize(Variable(O[idx,4:6,:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
